/*
 * simple_gm.h
 *
 *  Created on: Sep 9, 2011
 *      Author: tqlong
 */

#ifndef SIMPLE_GM_H_
#define SIMPLE_GM_H_

#include <map>
#include <armadillo>

typedef int Node;

/** Undirected edge
 *     i_ < j_
 */
struct Edge
{
  int i_;
  int j_;
  Edge(int i, int j, bool unique = false);

  struct EdgeCompare
  {
    bool operator()(const Edge& lhs, const Edge& rhs) const;
  };

  int other_end(int i) const
  {
    return (i == i_) ? j_ : i_;
  }
};

/** Simple graphical models consist of
 *    + a graph with node and edges
 *    + node potentials and edge potentials
 *    + assignment to node is binary
 *  Potentials are tables whose entries correspond to assignment to the node or the edges' nodes
 *  P(X) = 1/Z exp(\sum_i \theta_i(x_i) + \sum_ij \theta_ij(x_i, x_j) )
 */
class SimpleGM
{
public:
  typedef std::map<Node, arma::vec>                    n_potential_type;
  typedef std::map<Edge, arma::mat, Edge::EdgeCompare> e_potential_type;
  typedef std::map<int, int>                           imap_type;
  typedef imap_type                                    assignment_type;
  typedef std::vector<Edge>                            evec_type;
  typedef std::map<Node, evec_type>                    adjacent_edge_type;

  typedef const std::pair<const Node, arma::vec>&      const_pair_node_t;
  typedef const std::pair<const Edge, arma::mat>&      const_pair_edge_t;
  typedef std::pair<const Node, arma::vec>&            pair_node_t;
  typedef std::pair<const Edge, arma::mat>&            pair_edge_t;
public:
  SimpleGM() {}
  /** Network manipulation */
  void load(const char* fileName);  //TODO
  /** Add a node with cardinality and potential */
  void addPotential(int i, const arma::vec& v);
  /** Add an edge with corresponding potentials */
  void addPotential(int i, int j, const arma::mat& v);
  /** Add node */
  int addNode(int cardinality);
  Edge addEdge(int i, int j);

  /** Computation */
  /** compute P_hat(X) = exp(\sum_i \theta_i(x_i) + \sum_ij \theta_ij(x_i, x_j) ) */
  double pseudo_probability(const imap_type& assignment, bool isLog = true);
  /** Compute Z = sum_x P_hat(x) \forall x */
  double partition_function(bool isLog = true);
  /** Compute all node and edge marginals */
  void naive_marginals();
  /** Compute MAP assignment */
  assignment_type naive_map();

  /** Evidence processing */
  /** Create a new graphical model based on evidence */
  SimpleGM reduce(const assignment_type& a);

  int n() const { return node_potentials_.size(); }
  int ne() const { return edge_potentials_.size(); }
public:
  /** Cardinality of domain of each node */
  imap_type cardinality_;
  /** Node potentials */
  n_potential_type node_potentials_;
  /** Edge potentials */
  e_potential_type edge_potentials_;
  /** Adjacent edges for all nodes */
  adjacent_edge_type adj_edges_;
  /** Node marginals */
  n_potential_type node_marginals_;
  /** Edge marginals */
  e_potential_type edge_marginals_;
  /** Information of original graphical model */
  imap_type old2new_, new2old_;

  /** Print info */
  void print_params() const;
  void print_marginals() const;
protected:
  /** Check complete assignment */
  bool isCompleteAssignment(const imap_type& assignment) const;
  double nodePotential(int i, int value) const;
  /** Init marginals to zeros */
  void init_marginals();
  bool next_assignment(std::vector<int>& vi);
  double edge_potentials(int i, int j, int vi, int vj) const;
};

double logadd(double a, double b);
double logsum(const arma::vec& v);
double logsum(const arma::mat& v);

#endif /* SIMPLE_GM_H_ */
